import torch
import cv2
import numpy as np
from algorithms.detect.yolov7.models.experimental import attempt_load
from algorithms.detect.yolov7.utils.general import non_max_suppression, scale_coords
from algorithms.detect.yolov7.utils.datasets import letterbox
from algorithms.detect.inferencer_base_detect import InferencerBaseDetect


class YoloV7Inferencer(InferencerBaseDetect):
    def _load_model(self, model_name):
        """加载 YOLOv7 模型"""
        print(f"Loading YOLOv7 model: {model_name}")
        model = attempt_load(model_name, map_location=self.device)  # 加载权重
        model.eval()
        return model

    def preprocess(self, image_path):
        """对图像进行预处理：resize、归一化、填充等"""
        img0 = cv2.imread(image_path)
        img = letterbox(img0, new_shape=640)[0]  # resize + padding
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.float() / 255.0  # 归一化
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        return img

    def postprocess(self, outputs):
        """NMS 后处理，提取检测结果"""
        pred = non_max_suppression(outputs, conf_thres=0.25, iou_thres=0.45)
        det = pred[0]
        results = []
        if len(det):
            for *xyxy, conf, cls in det:
                results.append({
                    "bbox": [x.item() for x in xyxy],
                    "score": conf.item(),
                    "class_id": int(cls.item())
                })
        return results

    def detect_image(self, image_path):
        """对单张图像进行检测"""
        img_tensor = self.preprocess(image_path)
        with torch.no_grad():
            output = self.model(img_tensor)[0]  # YOLOv7 输出是 list，取第一个元素
        result = self.postprocess(output)
        return result

    def visualize(self, image_path, detection_result):
        """绘制检测框和标签"""
        image = cv2.imread(image_path)
        for det in detection_result:
            x1, y1, x2, y2 = map(int, det["bbox"])
            label = f"{det['class_id']} - {det['score']:.2f}"
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        cv2.imshow("YOLOv7 Detection", image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()